package com.sys.ai.workflow;

import com.agentsflex.core.document.Document;
import com.agentsflex.core.llm.BaseLlm;
import com.agentsflex.core.llm.ChatOptions;
import com.agentsflex.core.llm.LlmConfig;
import com.agentsflex.core.llm.StreamResponseListener;
import com.agentsflex.core.llm.client.BaseLlmClientListener;
import com.agentsflex.core.llm.client.HttpClient;
import com.agentsflex.core.llm.client.LlmClient;
import com.agentsflex.core.llm.client.LlmClientListener;
import com.agentsflex.core.llm.client.impl.SseClient;
import com.agentsflex.core.llm.embedding.EmbeddingOptions;
import com.agentsflex.core.llm.response.AiMessageResponse;
import com.agentsflex.core.message.AiMessage;
import com.agentsflex.core.message.Message;
import com.agentsflex.core.parser.AiMessageParser;
import com.agentsflex.core.parser.impl.DefaultAiMessageParser;
import com.agentsflex.core.prompt.DefaultPromptFormat;
import com.agentsflex.core.prompt.Prompt;
import com.agentsflex.core.prompt.PromptFormat;
import com.agentsflex.core.store.VectorData;
import com.agentsflex.core.util.CollectionUtil;
import com.agentsflex.core.util.LogUtil;
import com.agentsflex.core.util.Maps;
import com.agentsflex.core.util.StringUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONPath;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;

/**
 * @author LGY
 * Create by 2025/6/4 16:02
 */
public class CustomerLLM extends BaseLlm<CustomerLLMConfig> {

    private static final PromptFormat promptFormat = new DefaultPromptFormat();

    private final Map<String, String> headers = new HashMap();
    private final HttpClient httpClient = new HttpClient();
    private final AiMessageParser aiMessageParser = DefaultAiMessageParser.getChatGPTMessageParser(false);
    private final AiMessageParser streamMessageParser = DefaultAiMessageParser.getChatGPTMessageParser(true);

    public CustomerLLM(CustomerLLMConfig config) {
        super(config);
        this.headers.put("Content-Type", "application/json");
        this.headers.put("Accept", "application/json");
        this.headers.put("Authorization", "Bearer " + (this.getConfig()).getApiKey());
    }

    public AiMessageResponse chat(Prompt prompt, ChatOptions options) {
        Consumer<Map<String, String>> headersConfig = (this.config).getHeadersConfig();
        if (headersConfig != null) {
            headersConfig.accept(this.headers);
        }

        String payload = promptToPayload(prompt, this.config, options, false);
        String endpoint = (this.config).getEndpoint();
        String response = this.httpClient.post(endpoint + "/chat/completions", this.headers, payload);
        if ((this.config).isDebug()) {
            LogUtil.println(">>>>receive payload:" + response);
        }

        if (StringUtil.noText(response)) {
            return AiMessageResponse.error(prompt, response, "no content for response.");
        } else if (response.startsWith("{")) {
            JSONObject jsonObject = JSON.parseObject(response);
            Integer code = jsonObject.getInteger("code");
            return code != null ? AiMessageResponse.error(prompt, response, jsonObject.getString("message")) : new AiMessageResponse(prompt, response, (AiMessage) this.aiMessageParser.parse(jsonObject));
        } else {
            return AiMessageResponse.error(prompt, response, response);
        }
    }

    public void chatStream(Prompt prompt, StreamResponseListener streamResponseListener, ChatOptions chatOptions) {
        LlmClient llmClient = new SseClient();
        String payload = promptToPayload(prompt, this.config, chatOptions, true);
        String endpoint = (this.config).getEndpoint();
        LlmClientListener clientListener = new BaseLlmClientListener(this, llmClient, streamResponseListener, prompt, this.streamMessageParser);
        llmClient.start(endpoint + "/chat/completions", this.headers, payload, clientListener, this.config);
    }

    public void chatStream(Prompt prompt, StreamResponseListener streamResponseListener) {
        this.chatStream(prompt, streamResponseListener, ChatOptions.DEFAULT);
    }

    public VectorData embed(Document document, EmbeddingOptions embeddingOptions) {
        Consumer<Map<String, String>> headersConfig = (this.config).getHeadersConfig();
        if (headersConfig != null) {
            headersConfig.accept(this.headers);
        }
        String payload = Maps.of("model", Optional.ofNullable(embeddingOptions.getModel())
                        .orElse(config.getDefaultEmbeddingModel()))
                .set("input", document.getContent())
                .set("encoding_format", Optional.ofNullable(embeddingOptions.getEncodingFormat())
                        .orElse("float")).toJSON();
        String endpoint = (this.config).getEndpoint();
        String response = this.httpClient.post(endpoint + "/embeddings", this.headers, payload);
        if ((this.config).isDebug()) {
            LogUtil.println(">>>>receive payload:" + response);
        }

        if (StringUtil.noText(response)) {
            return null;
        } else if (response.startsWith("{")) {
            JSONObject jsonObject = JSON.parseObject(response);
            Integer code = jsonObject.getInteger("code");
            if (code != null) {
                return null;
            } else {
                VectorData vectorData = new VectorData();
                vectorData.setVector((double[]) JSONPath.read(response, "$.data[0].embedding"));
                return vectorData;
            }
        } else {
            return null;
        }
    }

    public static String promptToPayload(Prompt prompt, LlmConfig config, ChatOptions options, boolean withStream) {
        List<Message> messages = prompt.toMessages();
        Message message = CollectionUtil.lastItem(messages);
        return Maps.of("model", Optional.ofNullable(options.getModel())
                        .orElse(config.getModel()))
                .set("messages", promptFormat.toMessagesJsonObject(messages))
                .setIf(withStream, "stream", true)
                .setIfNotEmpty("tools", promptFormat.toFunctionsJsonObject(message))
                .setIfNotNull("top_p", options.getTopP())
                .setIfNotNull("top_k", options.getTopK())
                .setIfNotEmpty("stop", options.getStop())
                .setIf((map) -> {
                    return !map.containsKey("tools") && options.getTemperature() > 0.0F;
                }, "temperature", options.getTemperature()).setIf((map) -> {
                    return !map.containsKey("tools") && options.getMaxTokens() != null;
                }, "max_tokens", options.getMaxTokens()).toJSON();
    }
}
